# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Finetuning the library models for sequence classification on GLUE (Bert, XLM, XLNet, RoBERTa)."""

# 只用pheme数据集，然后用train_test_split方式划分数据集

from __future__ import absolute_import, division, print_function

import argparse
import glob
import logging
import os
import random
import torch
import csv
import copy
import numpy as np
import pandas as pd
from torch import nn

from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler, Dataset, TensorDataset)
from torch.utils.data.distributed import DistributedSampler
from transformers import DataProcessor, InputExample, InputFeatures
from transformers.modeling_utils import PreTrainedModel, SequenceSummary
from zmq import device

try:
    from torch.utils.tensorboard import SummaryWriter
except:
    from tensorboardX import SummaryWriter

from tqdm import tqdm, trange

from transformers import (WEIGHTS_NAME, BertConfig, BertModel, BertTokenizer,
                          RobertaConfig, RobertaModel, RobertaTokenizer,
                          XLNetConfig, XLNetModel, XLNetTokenizer,
                          DistilBertConfig, DistilBertModel, DistilBertTokenizer,
                          AlbertConfig, AlbertModel, AlbertTokenizer,
                          )
                          
from transformers import AdamW,  get_linear_schedule_with_warmup
from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_output_modes as output_modes
from transformers import glue_processors as processors
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP
from transformers import ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
from transformers import DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
from transformers import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP
from transformers import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP
from transformers import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP
from sklearn.metrics import f1_score
from sklearn.model_selection import train_test_split

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '5678'
os.environ["RANK"] = "0"        #进程序号  #local_rank是进程内的gpu编号
os.environ['WORLD_SIZE'] = '1'  #进程总数
os.environ["CUDA_VISIBLE_DEVICES"] = '1, 2'


logger = logging.getLogger(__name__)

#ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) for conf in (BertConfig, XLNetConfig, RobertaConfig, DistilBertConfig, AlbertConfig)), ())

#ALL_MODELS=tuple(BERT_PRETRAINED_CONFIG_ARCHIVE_MAP)
ALL_MODELS=tuple(ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP)
#ALL_MODELS=tuple(XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP)
#ALL_MODELS=tuple(DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP)
#ALL_MODELS=tuple(ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP)

MODEL_CLASSES = {
    'bert': (BertConfig, BertModel, BertTokenizer),
    'xlnet': (XLNetConfig, XLNetModel, XLNetTokenizer),
    'roberta': (RobertaConfig, RobertaModel, RobertaTokenizer),
    'distilbert': (DistilBertConfig, DistilBertModel, DistilBertTokenizer),
    'albert': (AlbertConfig, AlbertModel, AlbertTokenizer)
}


def set_seed(args):
    random.seed(args.seed)     # 设置随机数生成器的种子
    np.random.seed(args.seed)  # 用于生成指定随机数，“堆”
    torch.manual_seed(args.seed)  #设置随机数种子，保证每次运行文件生成的随机数是一样的，返回torch.Generator对象
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)   # 为所有的GPU设置种子


def simple_accuracy(preds, labels):
    return (preds == labels).mean()  # mean函数是求平均值，axis=0——压缩行，axis=1——压缩列


def acc_and_f1(preds, labels):
    acc = simple_accuracy(preds, labels)
    f1 = f1_score(y_true=labels, y_pred=preds,average='weighted')  # F1分数可以被解释为精度和召回率的加权平均。
    return {
            "acc": acc,
            "f1": f1,
            "acc_and_f1": (acc + f1) / 2,
    }


class TextDataset(Dataset):
    
    def __init__(self, data_list):
        self.data_list = data_list
        
    def __getitem__(self, idx):
        return (self.data_list[idx]["source_tweet_text"], self.data_list[idx]["replies_tweet_text_list"], self.data_list[idx]["label_b"])
    
    def __len__(self):
        return len(self.data_list)


fixed_replies_num = 16

class CnesProcessor(DataProcessor):
    """Processor for the cnews data set (GLUE version)."""

    def get_example_from_tensor_dict(self, tensor_dict):
        """See base class."""
        return InputExample(tensor_dict['idx'].numpy(),
                            tensor_dict['sentence'].numpy().decode('utf-8'),
                            None,
                            str(tensor_dict['label'].numpy()))

    def get_post_examples(self, data_dir):
        """See base class."""
        return self.create_post_examples(
            self._read_tsv(os.path.join(data_dir, "post.csv")), "post")
        
    def get_comment_examples(self, data_dir):
        """See base class."""
        return self.create_comment_examples(
            self._read_tsv(os.path.join(data_dir, "comment.csv")), "comment")

    def get_labels(self):
        """See base class."""
        return ["true",
                "false",
                "unverified"
                ]

    def create_post_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)
            text_a = line[1]
            label = line[0]
            examples.append(
                InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
        return examples
    
    def create_comment_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)
            text_a = line[0]
            examples.append(
                InputExample(guid=guid, text_a=text_a, text_b=None, label=None))
        return examples
        
    def source_data_change(self, source_data_dir):
        return self._data_change(os.path.join(source_data_dir, "train.csv"))
            #self._read_tsv(os.path.join(source_data_dir, "testing_features_extraction.csv")))
            
    def _data_change(self, source_data_dir):
        pd.set_option('display.float_format',lambda x : '%.3f' % x)  #禁用科学计数法
        all_dataframe = pd.read_csv(source_data_dir, lineterminator="\n")
        all_dataframe.rename(columns={"label_b\r": "label_b"}, inplace=True)
        all_dataframe["label_b"] = all_dataframe["label_b"].apply(lambda x: x.replace('\r', ''))
        all_dataframe["in_reply_to_status_id"].fillna(0, inplace=True)
        all_dataframe["in_reply_to_user_id_str"].fillna(0, inplace=True)
        all_dataframe["in_reply_to_status_id"] = all_dataframe["in_reply_to_status_id"].astype(int).astype(str)
        all_dataframe["in_reply_to_user_id_str"] = all_dataframe["in_reply_to_user_id_str"].astype(int).astype(str)
        all_dataframe[["belong_to_which_source_tweet", "twitter_id", "user_id"]] = all_dataframe[["belong_to_which_source_tweet", "twitter_id", "user_id"]].astype(str)
        all_dataframe["text"].fillna("null", inplace=True)
        remove_duplicated_df = all_dataframe.drop_duplicates(subset=["belong_to_which_source_tweet", "twitter_id"], keep="first").reset_index(drop=True)   #去掉重复项
        
        belong_to_which_source_tweet_map = {}
        for index, row in remove_duplicated_df.iterrows():
            item_dict = {
                "belong_to_which_source_tweet": row["belong_to_which_source_tweet"],
                "twitter_id": row["twitter_id"],
                "in_reply_to_status_id": row["in_reply_to_status_id"],
                "text": row["text"],
                "favorite_count": row["favorite_count"],
                "retweeted": row["retweeted"],
                "retweet_count": row["retweet_count"],
                "user_profile": {
                    "id": row["user_id"],
                    "verified": row["user_verified"],
                    "followers_count": row["user_followers_count"],
                    "listed_count": row["user_listed_count"],
                    "statuses_count": row["user_statuses_count"],
                    "description": row["user_description"],
                    "friends_count": row["user_friends_count"],
                    "location": row["user_location"],
                    "following": row["user_following"],
                    "geo_enabled": row["user_geo_enabled"],
                    "name": row["user_name"],
                    "favourites_count": row["user_favourites_count"],
                    "screen_name": row["user_screen_name"],
                    "created_at": row["user_created_at"],
                    "time_zone": row["user_time_zone"]
                },
                "in_reply_to_user_id_str": row["in_reply_to_user_id_str"],
                "created_at": row["created_at"],
                "place": row["place"],
                "label_a": row["label_a"],
                "label_b": row["label_b"]
            }
            if row["belong_to_which_source_tweet"] not in belong_to_which_source_tweet_map:
                belong_to_which_source_tweet_map[row["belong_to_which_source_tweet"]] = {
                    "label": row["label_b"],
                    "item_dict_list": []
                }
                belong_to_which_source_tweet_map.get(row["belong_to_which_source_tweet"])["item_dict_list"].append(item_dict)
            else:
                belong_to_which_source_tweet_map.get(row["belong_to_which_source_tweet"])["item_dict_list"].append(item_dict)
                
        
        data_list = []
        for key in belong_to_which_source_tweet_map:
            source_tweet_id = None
            replies_tweet_id_list = []
            source_tweet_text = None
            replies_tweet_text_list = []
            source_tweet_user_profile = {}
            replies_tweet_user_profile_list = []
            source_tweet_label_a = None
            replies_label_a_list = []
            label_b = belong_to_which_source_tweet_map.get(key)["label"]
            for item_dict in belong_to_which_source_tweet_map.get(key)["item_dict_list"]:
                if(item_dict["belong_to_which_source_tweet"] == item_dict["twitter_id"]):
                    source_tweet_id = str(item_dict["twitter_id"])
                    source_tweet_text = item_dict["text"]
                    source_tweet_user_profile = {
                        "id": str(item_dict["user_profile"]["id"]),
                        "verified": item_dict["user_profile"]["verified"],
                        "followers_count": item_dict["user_profile"]["followers_count"],
                        "listed_count": item_dict["user_profile"]["listed_count"],
                        "statuses_count": item_dict["user_profile"]["statuses_count"],
                        "description": item_dict["user_profile"]["description"],
                        "friends_count": item_dict["user_profile"]["friends_count"],
                        "location": item_dict["user_profile"]["location"],
                        "following": item_dict["user_profile"]["following"],
                        "geo_enabled": item_dict["user_profile"]["geo_enabled"],
                        "name": item_dict["user_profile"]["name"],
                        "favourites_count": item_dict["user_profile"]["favourites_count"],
                        "screen_name": item_dict["user_profile"]["screen_name"],
                        "created_at": item_dict["user_profile"]["created_at"],
                        "time_zone": item_dict["user_profile"]["time_zone"]
                    }
                    source_tweet_label_a = item_dict["label_a"]
                else:
                    replies_tweet_id_list.append(str(item_dict["twitter_id"]))
                    replies_tweet_text_list.append(item_dict["text"])
                    replies_tweet_user_profile_list.append({
                        "id": str(item_dict["user_profile"]["id"]),
                        "verified": item_dict["user_profile"]["verified"],
                        "followers_count": item_dict["user_profile"]["followers_count"],
                        "listed_count": item_dict["user_profile"]["listed_count"],
                        "statuses_count": item_dict["user_profile"]["statuses_count"],
                        "description": item_dict["user_profile"]["description"],
                        "friends_count": item_dict["user_profile"]["friends_count"],
                        "location": item_dict["user_profile"]["location"],
                        "following": item_dict["user_profile"]["following"],
                        "geo_enabled": item_dict["user_profile"]["geo_enabled"],
                        "name": item_dict["user_profile"]["name"],
                        "favourites_count": item_dict["user_profile"]["favourites_count"],
                        "screen_name": item_dict["user_profile"]["screen_name"],
                        "created_at": item_dict["user_profile"]["created_at"],
                        "time_zone": item_dict["user_profile"]["time_zone"]
                    })
                    replies_label_a_list.append(item_dict["label_a"])
            data_list.append({
                "source_tweet_id": source_tweet_id,
                "replies_tweet_id_list": replies_tweet_id_list,
                "source_tweet_text": source_tweet_text,
                "replies_tweet_text_list": replies_tweet_text_list,
                "source_tweet_user_profile": source_tweet_user_profile,
                "replies_tweet_user_profile_list": replies_tweet_user_profile_list,
                "source_tweet_label_a": source_tweet_label_a,
                "replies_label_a_list": replies_label_a_list,
                "label_b": label_b
            })
    
        
        for elem in data_list:
            if len(elem["replies_tweet_id_list"]) < fixed_replies_num:
                for i in range(fixed_replies_num - len(elem["replies_tweet_id_list"])):
                    elem["replies_tweet_id_list"].append("0")
                    elem["replies_tweet_text_list"].append("")
                    elem["replies_tweet_user_profile_list"].append("null")
                    elem["replies_label_a_list"].append("")
            elif len(elem["replies_tweet_id_list"]) > fixed_replies_num:
                index_to_renew_list = []
                index_to_renew_list = np.random.choice(len(elem["replies_tweet_id_list"]) - 1, fixed_replies_num, replace=False)
                new_replies_tweet_id_list = copy.deepcopy(elem["replies_tweet_id_list"])
                new_replies_tweet_text_list = copy.deepcopy(elem["replies_tweet_text_list"])
                new_replies_tweet_user_profile_list = copy.deepcopy(elem["replies_tweet_user_profile_list"])
                new_replies_label_a_list = copy.deepcopy(elem["replies_label_a_list"])
                elem["replies_tweet_id_list"] = []
                elem["replies_tweet_text_list"] = []
                elem["replies_tweet_user_profile_list"] = []
                elem["replies_label_a_list"] = []
                for index_to_renew in index_to_renew_list:
                    elem["replies_tweet_id_list"].append(new_replies_tweet_id_list[index_to_renew])
                    elem["replies_tweet_text_list"].append(new_replies_tweet_text_list[index_to_renew])
                    elem["replies_tweet_user_profile_list"].append(new_replies_tweet_user_profile_list[index_to_renew])
                    elem["replies_label_a_list"].append(new_replies_label_a_list[index_to_renew])
        
        dataset = TextDataset(data_list)
        
        target_data_dir = "./data"
        filepath1 = target_data_dir + "/post.csv" 
        filepath2 = target_data_dir + "/comment.csv" 

        for elem in dataset:
            item1 = {
                "source_tweet_text": elem[0],
                "label_b": elem[2]
                }
            with open(filepath1, 'a', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=["label_b", "source_tweet_text"], delimiter='\t')
                writer.writerow(item1)
            with open(filepath2, 'a', newline='') as f:
                writer = csv.writer(f, delimiter=',')
                for data in elem[1]:
                    writer.writerow([data])
                    
        return target_data_dir            
        
        
        '''target_data_dir = "./rumours"
        if set_type == "dev":
            filepath = target_data_dir + "/testing.csv" 
        else:
            filepath = target_data_dir + "/training.csv"
        for elem in data_list:
            item = {
                "source_tweet_text": elem["source_tweet_text"],
                "label_b": elem["label_b"]
            }
            with open(filepath, 'a', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=["label_b", "source_tweet_text"], delimiter='\t')
                writer.writerow(item) 
                
        #pd_all = pd.read_csv(target_data_dir + "/training.tsv", delimiter="\t", encoding='utf-8')
        #pd_all.to_csv(target_data_dir + "/training.csv", index=False, sep='\t', encoding='utf-8')
        return target_data_dir'''


processors["cnews"] = CnesProcessor
output_modes["cnews"] = "classification"


class TextClassifier(PreTrainedModel):

    def __init__(self, config, model_class, num_labels, model_name_or_path):
        super(TextClassifier, self).__init__(config)
        self.num_labels = num_labels
        self.model = model_class.from_pretrained(model_name_or_path)
        
        if self.model.base_model_prefix == 'bert':
            self.dim = config.hidden_size
            self.dropout_prob = config.hidden_dropout_prob
        elif self.model.base_model_prefix == 'distilbert':
            self.dim = config.dim
            self.dropout_prob = config.seq_classif_dropout
        elif self.model.base_model_prefix == "roberta":
            self.dropout_prob = config.hidden_dropout_prob
            self.dim = config.hidden_size
        elif self.model.base_model_prefix == "transformer":
            self.dim = config.d_model
            #self.sequence_summary = SequenceSummary(config)
            self.dropout_prob = config.dropout
        elif self.model.base_model_prefix == "albert":
            self.dropout_prob = config.classifier_dropout_prob
            self.dim = config.hidden_size
            
        self.pre_classifier = nn.Linear(self.dim*2, self.dim*2)
        self.classifier = nn.Linear(self.dim*2, self.num_labels)
        self.dropout = nn.Dropout(self.dropout_prob)             
        
    def forward(self, post_input_ids, post_attention_mask, post_labels, comment_input_ids, comment_attention_mask):
        
        post_outputs = self.model(input_ids = post_input_ids,
                            attention_mask = post_attention_mask,
                            )
        
        bs, rc, sq = comment_input_ids.shape[0], comment_input_ids.shape[1], comment_input_ids.shape[2]
        comment_input_ids = comment_input_ids.view(bs*rc, -1)
        comment_attention_mask = comment_attention_mask.view(bs*rc, -1)
        comment_outputs = self.model(input_ids = comment_input_ids,
                              attention_mask = comment_attention_mask)
        
        post_last_hidden_state = post_outputs[0]         # (bs, seq_len, dim)
        comment_last_hidden_state = comment_outputs[0]   # (bs*16, seq_len, dim)
        
        post_pooled_output = post_last_hidden_state[:, 0]             # (bs, dim)
        comment_pooled_output = comment_last_hidden_state[:, 0]       # (bs*16, dim)
        
        dim= post_pooled_output.shape[1]
        comment_pooled_output = comment_pooled_output.view(bs, -1, dim)   # (bs, 16, dim)
        comment_pooled_output = torch.mean(comment_pooled_output, axis=1)    #(bs, dim)
        
        pooled_output = torch.cat((post_pooled_output, comment_pooled_output), dim=1)   #(bs, dim*2)
        
        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim*2)
        pooled_output = nn.ReLU()(pooled_output)            # (bs, dim*2)
        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim*2)
        
        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim*2)
        pooled_output = nn.ReLU()(pooled_output)            # (bs, dim*2)
        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim*2)
        
        dropout_output = self.dropout(pooled_output)        # (bs, dim*2)
        logits = self.classifier(dropout_output)            # (bs, num_labels)
        
        if self.model.base_model_prefix == 'distilbert':
            outputs = (logits,) + post_outputs[1:] + comment_outputs[1:]
        else:
            outputs = (logits,) + post_outputs[2:] + comment_outputs[2:]
        
        if post_labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), post_labels.view(-1))
            outputs = (loss,) + outputs
            #outputs = (loss,)

        return outputs
        

def train(args, train_dataset, eval_dataset, model, tokenizer):
    """ Train the model """
    if args.local_rank in [-1, 0, 1]:
        tb_writer = SummaryWriter('./runs/bert')

    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1  #//表示整数除法，向下取整
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]

    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps,
                                                num_training_steps=t_total)
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)
    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
                args.train_batch_size * args.gradient_accumulation_steps * (
                    torch.distributed.get_world_size() if args.local_rank != -1 else 1))
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch")
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)
    for epoch_num in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration")
        train_loss = 0.0
        train_acc = 0.0
        train_step = 0.0
        for step, batch in enumerate(epoch_iterator):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            
            post_inputs = {'input_ids': batch[0],
                      'attention_mask': batch[1],
                      'labels': batch[2]}
            comment_inputs = {'input_ids': batch[3],
                      'attention_mask': batch[4]}
            
            outputs = model(post_inputs['input_ids'], post_inputs['attention_mask'], post_inputs['labels'], 
                            comment_inputs["input_ids"], comment_inputs["attention_mask"])
            
            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)

            if args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            train_loss += loss.item()
            
            acc = (outputs[1].argmax(dim=1) == post_inputs['labels']).sum().item()
            train_acc += acc
            
            if (step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1
                train_step += 1

                if args.local_rank in [-1, 0, 1] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    # Log metrics
                    if args.local_rank == -1 and args.evaluate_during_training:  # Only evaluate when single GPU otherwise metrics may not average well
                        results = evaluate(args, eval_dataset, model, tokenizer)
                        for key, value in results.items():
                            tb_writer.add_scalar('eval_{}'.format(key), value, global_step)
                    tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
                    tb_writer.add_scalar('loss', (tr_loss - logging_loss) / args.logging_steps, global_step)
                    logging_loss = tr_loss

                if args.local_rank in [-1, 0, 1] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    # Save model checkpoint
                    output_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    model_to_save = model.module if hasattr(model,
                                                            'module') else model  # Take care of distributed/parallel training
                    model_to_save.save_pretrained(output_dir)
                    torch.save(args, os.path.join(output_dir, 'training_args.bin'))
                    logger.info("Saving model checkpoint to %s", output_dir)
                
                del post_inputs
                del comment_inputs
                torch.cuda.empty_cache()

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        
        print(
            f'Epochs: {epoch_num + 1} | Train Loss: {train_loss / train_step: .3f} | Train Accuracy: {train_acc / len(train_dataset): .3f} ')
        
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

    if args.local_rank in [-1, 0, 1]:
        tb_writer.close()

    return global_step, tr_loss / global_step


def evaluate(args, eval_dataset, model, tokenizer, prefix=""):
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_task_names = ("mnli", "mnli-mm") if args.task_name == "mnli" else (args.task_name,)
    eval_outputs_dirs = (args.output_dir, args.output_dir + '-MM') if args.task_name == "mnli" else (args.output_dir,)

    results = {}
    for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):
        eval_dataset = eval_dataset

        if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0, 1]:
            os.makedirs(eval_output_dir)

        args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
        
        eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
        eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, shuffle=False)        

        # multi-gpu eval
        if args.n_gpu > 1:
            model = torch.nn.DataParallel(model)

        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(eval_dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)
        eval_loss = 0.0
        nb_eval_steps = 0
        preds = None
        out_label_ids = None
        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.eval()
            batch = tuple(t.to(args.device) for t in batch)

            with torch.no_grad():
                post_inputs = {'input_ids': batch[0],
                      'attention_mask': batch[1],
                      'labels': batch[2]}
                comment_inputs = {'input_ids': batch[3],
                      'attention_mask': batch[4]}
                outputs = model(post_inputs['input_ids'], post_inputs['attention_mask'], post_inputs['labels'], 
                                comment_inputs["input_ids"], comment_inputs["attention_mask"])
                tmp_eval_loss, logits = outputs[:2]

                eval_loss += tmp_eval_loss.mean().item()
            nb_eval_steps += 1
            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = post_inputs['labels'].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(out_label_ids, post_inputs['labels'].detach().cpu().numpy(), axis=0)

        eval_loss = eval_loss / nb_eval_steps / args.eval_batch_size
        if args.output_mode == "classification":
            preds = np.argmax(preds, axis=1)
        elif args.output_mode == "regression":
            preds = np.squeeze(preds)

        if eval_task == "cnews":
            result = acc_and_f1(preds,out_label_ids)
        else:
            result = compute_metrics(eval_task, preds, out_label_ids)
        results.update(result)

        output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
        with open(output_eval_file, "w") as writer:
            logger.info("***** Eval results {} *****".format(prefix))
            logger.info("  eval_loss = %s", str(eval_loss))
            for key in sorted(result.keys()):
                logger.info("  %s = %s", key, str(result[key]))
                writer.write("%s = %s\n" % (key, str(result[key])))

    return results

def predict(args, pred_dataset, model, tokenizer, prefix=""):
    
    pred_dataset = pred_dataset

    if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
        os.makedirs(args.output_dir)
    #测试阶段batch大小
    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    #定义采样方式
    eval_sampler = SequentialSampler(pred_dataset) if args.local_rank == -1 else DistributedSampler(pred_dataset)
    #测试集Dataloader
    eval_dataloader = DataLoader(pred_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

    # multi-gpu eval 单机多卡
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Predict!
    logger.info("***** Running predict {} *****".format(prefix))
    logger.info("  Num examples = %d", len(pred_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
    eval_loss = 0.0
    nb_eval_steps = 0
    preds = None    #为预测值
    for batch in tqdm(eval_dataloader, desc="Predicting"):
        model.eval() #测试模式
        batch = tuple(t.to(args.device) for t in batch)#把测试数据转移到设备上

        with torch.no_grad():#关闭梯度计算
            #构建模型输入 字典形式。 token_type_ids为batch[2] 分类任务为单输入句子 默认全为0
            with torch.no_grad():
                post_inputs = {'input_ids': batch[0],
                      'attention_mask': batch[1],
                      'labels': batch[2]}
                comment_inputs = {'input_ids': batch[3],
                      'attention_mask': batch[4]}
                outputs = model(post_inputs['input_ids'], post_inputs['attention_mask'], post_inputs['labels'], 
                                comment_inputs["input_ids"], comment_inputs["attention_mask"])
            tmp_eval_loss, logits = outputs[:2] #前两项为loss、logits

            eval_loss += tmp_eval_loss.mean().item()
        nb_eval_steps += 1
        if preds is None:
            preds = logits.detach().cpu().numpy()
        else:
            preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)

    eval_loss = eval_loss / nb_eval_steps
    if args.output_mode == "classification":
        preds = np.argmax(preds, axis=1) #得到预测的标签 对预测结果按行取argmax

    #把预测的标签 输出为csv文件
    pd.DataFrame(preds).to_csv(os.path.join(args.output_dir, "predicted.csv"), index=False, header=None)


def load_and_cache_examples(args, task, tokenizer, evaluate=False):
    if args.local_rank not in [-1, 0, 1] and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    processor = processors[task]()
    output_mode = output_modes[task]
    # Load data features from cache or dataset file
    cached_post_features_file = os.path.join(args.data_dir, 'cached_{}_{}_{}_{}'.format(
        'post',
        'distilbert',
        str(args.max_seq_length),
        str(task)))
    cached_comment_features_file = os.path.join(args.data_dir, 'cached_{}_{}_{}_{}'.format(
        'comment',
        'distilbert',
        str(args.max_seq_length),
        str(task)))
    if os.path.exists(cached_post_features_file) and os.path.exists(cached_comment_features_file) and not args.overwrite_cache:
        logger.info("Loading post_features from cached file %s", cached_post_features_file)
        post_features = torch.load(cached_post_features_file)
        logger.info("Loading comment_features from cached file %s", cached_post_features_file)
        comment_features = torch.load(cached_comment_features_file)
    else:
        logger.info("Creating features from dataset file at %s", args.data_dir)
        label_list = processor.get_labels()
        if task in ['mnli', 'mnli-mm'] and args.model_type in ['roberta']:
            # HACK(label indices are swapped in RoBERTa pretrained model)
            label_list[1], label_list[2] = label_list[2], label_list[1]
                   
        #插入对数据集的处理，使其变为label、text的形式
        data_dir = processor.source_data_change(args.data_dir)
        post_examples = processor.get_post_examples(data_dir)
        comment_examples = processor.get_comment_examples(data_dir)
        post_features = convert_examples_to_features(post_examples, tokenizer, label_list=label_list, 
                                                max_length=args.max_seq_length, output_mode=output_mode,
                                                )
        comment_features = convert_examples_to_features(comment_examples, tokenizer, label_list=label_list, 
                                                max_length=args.max_seq_length, output_mode=output_mode,
                                                )
        if args.local_rank in [-1, 0, 1]:
            logger.info("Saving post_features into cached file %s", cached_post_features_file)
            torch.save(post_features, cached_post_features_file)
            logger.info("Saving comment_features into cached file %s", cached_comment_features_file)
            torch.save(comment_features, cached_comment_features_file)

    if args.local_rank == 0 and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    # Convert to Tensors and build dataset
    all_post_input_ids = torch.tensor([f.input_ids for f in post_features], dtype=torch.long)
    all_post_attention_mask = torch.tensor([f.attention_mask for f in post_features], dtype=torch.long)
    #all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
    if output_mode == "classification":
        all_post_labels = torch.tensor([f.label for f in post_features], dtype=torch.long)
    elif output_mode == "regression":
        all_post_labels = torch.tensor([f.label for f in post_features], dtype=torch.float)
        
    all_comment_input_ids = torch.tensor([f.input_ids for f in comment_features], dtype=torch.long)
    all_comment_attention_mask = torch.tensor([f.attention_mask for f in comment_features], dtype=torch.long)
    
    #这里改变了！！！
    all_comment_input_ids_list = []
    for i in range(len(all_comment_input_ids)):
        if i % 16 == 0:
            comment_input_ids_list = []
            comment_input_ids_list.append(all_comment_input_ids[i])
        else:
            comment_input_ids_list.append(all_comment_input_ids[i])
        if len(comment_input_ids_list) == 16:
            comment_input_ids_list = torch.stack(comment_input_ids_list)
            all_comment_input_ids_list.append(comment_input_ids_list)
    all_comment_input_ids_list = torch.stack(all_comment_input_ids_list)
        
    all_comment_attention_mask_list = []
    for i in range(len(all_comment_attention_mask)):
        if i % 16 == 0:
            comment_attention_mask_list = []
            comment_attention_mask_list.append(all_comment_attention_mask[i])
        else:
            comment_attention_mask_list.append(all_comment_attention_mask[i])
        if len(comment_attention_mask_list) == 16:
            comment_attention_mask_list = torch.stack(comment_attention_mask_list)
            all_comment_attention_mask_list.append(comment_attention_mask_list)               
    all_comment_attention_mask_list = torch.stack(all_comment_attention_mask_list)   
    
    dataset = TensorDataset(all_post_input_ids, all_post_attention_mask, all_post_labels, 
                            all_comment_input_ids_list, all_comment_attention_mask_list)       
    
    return dataset


def main():
    parser = argparse.ArgumentParser()

    ## Required parameters
    parser.add_argument("--data_dir", default='./data', type=str, required=False,
                        help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
    parser.add_argument("--model_type", default='bert', type=str, required=False,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
    parser.add_argument("--model_name_or_path", default='bert-base-uncased', type=str, required=False,
                        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS))
    parser.add_argument("--task_name", default='cnews', type=str, required=False,
                        help="The name of the task to train selected in the list: " + ", ".join(processors.keys()))
    parser.add_argument("--output_dir", default='./outs/bert', type=str, required=False,
                        help="The output directory where the model predictions and checkpoints will be written.")

    ## Other parameters
    parser.add_argument("--config_name", default="", type=str,
                        help="Pretrained config name or path if not the same as model_name")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument("--cache_dir", default="", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")
    parser.add_argument("--max_seq_length", default=128, type=int,
                        help="The maximum total input sequence length after tokenization. Sequences longer "
                             "than this will be truncated, sequences shorter will be padded.")
    parser.add_argument("--do_train", default=True, action='store_true',
                        help="Whether to run training.")
    parser.add_argument("--do_eval",default=True, action='store_true',
                        help="Whether to run eval on the dev set.")
    parser.add_argument("--do_predict", default=True, action='store_true',
                        help="Whether to run predict on the test set.")
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Rul evaluation during training at each logging step.")
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")

    parser.add_argument("--per_gpu_train_batch_size", default=4, type=int,
                        help="Batch size per GPU/CPU for training.")
    parser.add_argument("--per_gpu_eval_batch_size", default=4, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument("--learning_rate", default=2e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument("--weight_decay", default=0.0, type=float,
                        help="Weight deay if we apply some.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
                        help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument("--num_train_epochs", default=5.0, type=float,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--max_steps", default=-1, type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
    parser.add_argument("--warmup_steps", default=0, type=int,
                        help="Linear warmup over warmup_steps.")

    parser.add_argument('--logging_steps', type=int, default=100,
                        help="Log every X updates steps.")
    parser.add_argument('--save_steps', type=int, default=384,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument("--eval_all_checkpoints",default=True, action='store_true',
                        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
    parser.add_argument("--no_cuda", action='store_true',
                        help="Avoid using CUDA when available")
    parser.add_argument('--overwrite_output_dir', default='./rumours_outs', action='store_true',
                        help="Overwrite the content of the output directory")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")

    parser.add_argument('--fp16', action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
    parser.add_argument("--local_rank", type=int, default=1,
                        help="For distributed training: local_rank")
    parser.add_argument('--server_ip', type=str, default='', help="For distant debugging.")
    parser.add_argument('--server_port', type=str, default='', help="For distant debugging.")
    args = parser.parse_args()

    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir:
        raise ValueError(
            "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
                args.output_dir))

    # Setup distant debugging if needed
    if args.server_ip and args.server_port:
        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
        import ptvsd
        print("Waiting for debugger attach")
        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
        ptvsd.wait_for_attach()

    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        #os.environ["CUDA_VISIBLE_DEVICES"] = "1"
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        #device = torch.device('cuda:1')
        torch.distributed.init_process_group(backend='nccl')
        args.n_gpu = 1
    args.device = device

    # Setup logging
    '''logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)'''
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO if args.local_rank in [-1, 0, 1, 2, 3, 4] else logging.WARN)
    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)

    # Set seed
    set_seed(args)

    # Prepare GLUE task
    args.task_name = args.task_name.lower()
    if args.task_name not in processors:
        raise ValueError("Task not found: %s" % (args.task_name))
    processor = processors[args.task_name]()
    args.output_mode = output_modes[args.task_name]
    label_list = processor.get_labels()
    num_labels = len(label_list)

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0, 1]:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                          num_labels=num_labels,
                                          finetuning_task=args.task_name,
                                          cache_dir=args.cache_dir if args.cache_dir else None)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
                                                do_lower_case=args.do_lower_case,
                                                cache_dir=args.cache_dir if args.cache_dir else None)
    model = TextClassifier(config, model_class, num_labels, args.model_name_or_path)

    if args.local_rank == 0:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

    model.to(args.device)

    logger.info("Training/evaluation parameters %s", args)
    
    dataset = load_and_cache_examples(args, args.task_name, tokenizer, evaluate=False)
    train_dataset, eval_dataset = train_test_split(dataset, train_size=0.8, random_state=42)
    train_dataset, test_dataset = train_test_split(train_dataset, train_size=0.8, random_state=42)

    # 将test数据的label写入csv文件中
    test_data_list = []
    for elem in test_dataset:
        a = elem[2].item()
        test_data_list.append(a)
    with open(args.output_dir+"/test_data.csv", 'a', newline="") as f:
            writer = csv.writer(f, delimiter='\n')
            writer.writerow(test_data_list)
            
    # Training
    if args.do_train:
        global_step, tr_loss = train(args, train_dataset, eval_dataset, model, tokenizer)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

    # Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
    if args.do_train:
        # Create output directory if needed
        if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)

        logger.info("Saving model checkpoint to %s", args.output_dir)
        # Save a trained model, configuration and tokenizer using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        model_to_save = model.module if hasattr(model,
                                                'module') else model  # Take care of distributed/parallel training
        
        model_to_save.save_pretrained(args.output_dir)
        #config.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))

        # Load a trained model and vocabulary that you have fine-tuned
        #model = torch.load(args.output_dir)
        '''config = config_class.from_pretrained(args.output_dir)
        tokenizer = tokenizer_class.from_pretrained(args.output_dir)
        model = TextClassifier(config, model_class, num_labels, args.output_dir)
        model.load_state_dict(torch.load(args.output_dir + '/pytorch_model.bin'))
        model.to(args.device)'''

    # Evaluation
    results = {}
    if args.do_eval and args.local_rank in [-1, 0, 1]:
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
        checkpoints = [args.output_dir]
        if args.eval_all_checkpoints:
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
            prefix = "checkpoint-" + str(checkpoint.split("-")[-1]) if checkpoint.find('checkpoint') != -1 else ""
            
            config = config_class.from_pretrained(checkpoint)
            model = TextClassifier(config, model_class, num_labels, checkpoint)
            model.load_state_dict(torch.load(checkpoint + '/pytorch_model.bin'))
            model.to(args.device)
            result = evaluate(args, eval_dataset, model, tokenizer, prefix=prefix)
            result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())
            results.update(result)
            
    #Predict
    print("\n++++++++++++++++++++++++++=\n")
    print("executing!\n")
    if args.do_predict and args.local_rank in [-1, 0, 1]:
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
        '''config = config_class.from_pretrained(args.output_dir)
        model = TextClassifier(config, model_class, num_labels, args.output_dir)
        model.load_state_dict(torch.load(args.output_dir + '/pytorch_model.bin'))
        model.to(args.device)
        logger.info("Predict...")
        result = predict(args, test_post_dataset, model, tokenizer, prefix=prefix)'''
        checkpoints = [args.output_dir]
        if args.eval_all_checkpoints:
            checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split('/')[-1] if checkpoint.find('checkpoint') != -1 else ""
            
            config = config_class.from_pretrained(checkpoint)
            model = TextClassifier(config, model_class, num_labels, checkpoint)
            model.load_state_dict(torch.load(checkpoint + '/pytorch_model.bin'))
            model.to(args.device)
            logger.info("Predict...")
            result = predict(args, test_dataset, model, tokenizer, prefix=prefix)


    return results


if __name__ == "__main__":
    main()
